package com.agg.persist.ast;

import com.agg.core.utils.Asserts;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.util.JdbcConstants;


import java.util.*;


public class TableVisitor extends MySqlASTVisitorAdapter {

    private List<String> tableNames= new ArrayList<>();

    public boolean visit(SQLExprTableSource sqlExprTableSource) {
        String taleName = sqlExprTableSource.getName().getSimpleName();
        tableNames.add(taleName); return true; }
        public List<String> getTableNames() {
        return tableNames;
    }
      public static void main2(String[] args) {
          //  String sqlStr = "SELECT A.ID,B.ID FROM A LEFT JOIN B ON A.ID =B.ID LEFT JOIN C ON B.ID = C.ID";

            String sqlStr="SELECT A.ID,B.ID FROM A LEFT JOIN B";
            String dbType = JdbcConstants.MYSQL;
            SQLStatement stmt = SQLUtils.parseSingleStatement(sqlStr, dbType);
            TableVisitor tableVisitor = new TableVisitor();
            stmt.accept(tableVisitor);
            System.out.println(tableVisitor.getTableNames());
      }


    public static void main(String[] args) {
        //  String sqlStr = "SELECT A.ID,B.ID FROM A LEFT JOIN B ON A.ID =B.ID LEFT JOIN C ON B.ID = C.ID";

        String sqlStr="SELECT A.ID,count(B.TT) cnt,B.ID as ID_2 FROM A LEFT JOIN (select ID, kk from  tt) as B ";
        String dbType = JdbcConstants.MYSQL;
        List<SQLStatement> stmts = SQLUtils.parseStatements(sqlStr, dbType);
        List<String> columnList = new ArrayList<>();

        for(SQLStatement statement:stmts){
            SchemaStatVisitor visitor = new SchemaStatVisitor(dbType);

            statement.accept(visitor);
            SQLSelectStatement selectStatement = (SQLSelectStatement) statement;
            SQLSelectQueryBlock queryBlock =  selectStatement.getSelect().getFirstQueryBlock();

            Collection<TableStat.Column> columns = visitor.getColumns();
            columns.stream().forEach(row -> {
                if(row.isSelect()){
                    //保存select字段
                    columnList.add(row.getName());
                }
            });
        }
        for(String str:columnList){
          //  System.out.println(str);
        }
        SQLStatement stmt = stmts.get(0);
        List<SQLSelectItem> list = ((MySqlSelectQueryBlock)((SQLSelect)((SQLSelectStatement)stmt).getSelect()).getQuery()).getSelectList();
        Set<String> set = new HashSet<>();
        for(SQLSelectItem item : list) {
            String selectColumnAlias = item.getAlias();
            String clomn=null;
            SQLExpr expr = item.getExpr();
            if (expr instanceof SQLPropertyExpr) {
                SQLPropertyExpr selectColumnExpr =(SQLPropertyExpr) expr;
            }else if(expr instanceof SQLAggregateExpr){
                SQLAggregateExpr aggregateExpr=(SQLAggregateExpr)expr;
                System.out.println(item.getAlias());
            }


            set.add(clomn);

        }

     //   getColumn(33L,sqlStr);
    }


    public static Set<String> getColumn(Long type,String sql){

        if(Asserts.isEmpty(sql)){
            return new HashSet<>();
        }
        List<SQLSelectItem> list = null;
        Set<String> set = new HashSet<>();
        if(type == 2){
            String dbType = JdbcConstants.HIVE;
            List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
            if(stmtList == null){
                return new HashSet<>();
            }
            SQLStatement stmt = stmtList.get(0);
            list = ((SQLSelectQueryBlock)((SQLSelect)((SQLSelectStatement)stmt).getSelect()).getQuery()).getSelectList();
        }else {
            String dbType = JdbcConstants.MYSQL;
            List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
            if(stmtList == null){
                return new HashSet<>();
            }
            SQLStatement stmt = stmtList.get(0);
            list = ((MySqlSelectQueryBlock)((SQLSelect)((SQLSelectStatement)stmt).getSelect()).getQuery()).getSelectList();
        }
        for(SQLSelectItem item : list){
            String clomn = item.getAlias();
            clomn = clomn.replaceAll("\"","");
            clomn = clomn.replaceAll("\'","");
            set.add(clomn);
        }

        return set;
    }

}
